import random

from utils.curl import curl_request, curl_request_file
import time
import os
import uuid
import traceback
import requests
import cv2
from utils import sr_s3
from typing import Optional
import http.client
import json
from dotenv import load_dotenv

load_dotenv(verbose=True, override=True)

api_key = os.getenv('PIXVERSE_API_KEY')

# 上传  cn  
UPLOAD_URL = 'https://app-api.pixverse.ai/openapi/v2/image/upload'
# 文生视频
TEXT_TO_VIDEO_URL = 'https://app-api.pixverse.ai/openapi/v2/video/text/generate'
# 图生视频
IMG_TO_VIDEO_URL = 'https://app-api.pixverse.ai/openapi/v2/video/img/generate'
# 首尾帧
FRAME_TO_VIDEO_URL = 'https://app-api.pixverse.ai/openapi/v2/video/transition/generate'
# 测试 sk-25ac79c560f96200a6cde49a32461808    正式 模板sk-ec3a6afc428eea05def5cd3338a182e7


def get_default_headers(trace_id=None):
    """获取默认请求头"""
    return {'Content-Type': 'application/json', 'API-KEY': api_key, 'Ai-trace-id': trace_id or str(uuid.uuid4())}


def get_video_info(video_path: str) -> dict:
    """获取视频的宽高、时长和帧率信息"""
    try:
        cap = cv2.VideoCapture(video_path)

        # 获取视频的基本信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # 计算视频时长（秒）
        duration = frame_count / fps if fps > 0 else 0

        cap.release()

        return {
            'width': width,
            'height': height,
            'duration': duration,
            'fps': fps
        }
    except Exception as e:
        print(f"获取视频信息失败：{str(e)}")
        # 返回默认值或部分信息
        return {
            'width': None,
            'height': None,
            'duration': None,
            'fps': None
        }


def download_video(video_url: str, save_dir: str = "./download", filename: Optional[str] = None) -> Optional[dict]:
    """
    从视频URL下载视频并保存到本地

    参数：
    - video_url: 视频下载URL
    - save_dir: 保存目录（默认 ./download）
    - filename: 自定义文件名（如不指定则从URL提取）

    返回：
    - 本地保存路径（成功）或 None（失败）
    """
    try:
        # 创建保存目录（如不存在）
        os.makedirs(save_dir, exist_ok=True)

        # 生成文件名（默认从URL提取）
        if not filename:
            # 从URL尾部提取文件名（如无则用随机名）
            filename = video_url.split("/")[-1]
            if not filename.endswith((".mp4", ".mov", ".avi")):
                filename = f"video_{os.urandom(4).hex()}.mp4"  # 随机名+默认mp4格式

        save_path = os.path.join(save_dir, filename)

        # 下载视频（带超时和进度提示）
        print(f"开始下载视频到 {save_path}...")
        response = requests.get(video_url, stream=True, timeout=30)
        response.raise_for_status()  # 检查HTTP错误

        # 分块写入文件（适合大文件）
        with open(save_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB/块
                if chunk:
                    f.write(chunk)
            # 获取视频宽高
        video_info = get_video_info(save_path)
        print(f"视频保存成功：{save_path}")

        return {
            'save_path': save_path,
            'width': video_info['width'],
            'height': video_info['height'],
            'duration': video_info['duration'],
            'fps': video_info['fps']
        }
        # return save_path
    except Exception as e:
        print(f"视频下载失败：{str(e)}")
        return None


def upload_video_to_s3(video_path, s3_storage, storage_service, bucket_name):
    """上传视频到R2并返回R2路径"""
    if not video_path or not os.path.exists(video_path):
        return None, "视频文件不存在"

    # try:
    # 调用S3工具上传视频
    s3_key = sr_s3.upload_login_or_visitor(
        video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
    )
    return s3_key, None
    # except Exception as e:
    #     return None, f"上传R2失败: {str(e)}"


# 图生视频
def create_task(data, url, trace_id=None):
    # time.sleep(5)
    headers = get_default_headers(trace_id)
    print("data:", data)
    print("url:", url)
    print("headers:", headers)
    response = curl_request('POST', url, data=data, headers=headers)
    print("POST响应:", response)
    if response["ErrCode"] == 0:
        print("响应成功")
        video_id = response["Resp"]["video_id"]
        # GET请求示例
        while True:
            # GET请求使用相同的headers，确保Ai-trace-id一致
            response = get_res(video_id, headers)
            print("GET响应:", response)
            if response["ErrCode"] == 0:
                if response["Resp"]["status"] == 5:
                    print("GET响应: 处理中")
                    time.sleep(3)
                else:
                    print("GET响应: 成功", response)
                    return response
            else:
                return {"ErrCode": response["ErrCode"], "ErrMsg": ""}
    else:
        print({"code": response["ErrCode"], "msg": response["ErrMsg"]})
        return {"ErrCode": response["ErrCode"], "ErrMsg": response["ErrMsg"]}


# 获取视频结果
def get_res(video_id, headers):
    try:
        conn = http.client.HTTPSConnection("app-api.pixverse.ai")
        payload = ''
        conn.request("GET", "/openapi/v2/video/result/{}".format(video_id), payload, headers)
        res = conn.getresponse()
        data = res.read().decode("utf-8")
        response = json.loads(data)
        return response
    except Exception as e:
        print("发生未知错误：", e)
        return {"ErrCode": 2, "ErrMsg": e}
    finally:
        # 确保连接关闭
        try:
            conn.close()
        except:
            pass

        # 假设的函数，用于从远程链接下载图片（实际实现）


def download_image(url, local_filename):
    try:
        response = requests.get(url)
        response.raise_for_status()  # 如果请求出错，引发HTTPError异常
        with open(local_filename, 'wb') as f:
            f.write(response.content)
        return True
    except Exception as err:
        print({"code": 404, "msg": err})
        return False


def upload_file(local_image_file, trace_id=None):
    headers = get_default_headers(trace_id)
    print("UPLOAD_URL", local_image_file)
    response = curl_request_file('POST', UPLOAD_URL, local_image_file, headers)
    # response = {'ErrCode': 0, 'ErrMsg': 'Success', 'Resp': {'img_id': 831292021, 'img_url': 'https://media.pixverse.ai/openapi%2Fd40f3129-a9d4-4afa-8ead-98a522f57f181d54fa67-bd1f-4ba5-a03a-69d4e931695b_auto.png'}}
    print("上传文件", response)
    return response


def download_single_image(img_url, storage_service, bucket_name, save_dir="./upload"):
    """下载单张图片到本地并返回路径"""
    try:
        print("下载单张图片到本地并返回路径", img_url)
        os.makedirs(save_dir, exist_ok=True)
        name = str(uuid.uuid4())
        ext = img_url[img_url.rfind('.'):] if '.' in img_url else '.jpg'
        save_path = os.path.join(save_dir, f"{name}{ext}")

        # 调用S3工具下载图片
        success, err = sr_s3.download_file(
            img_url, save_path, serviceName=storage_service, bucketName=bucket_name
        )
        if not success:
            return None, f"下载失败: {err}"
        # # JPG/PNG/WebP
        # valid_extensions = ('.jpg', '.jpeg', '.png', '.webp')
        # # 压缩图片（确保符合模型要求）
        # success, new_path = sr_utils.video_compress_image_size_path(save_path, 20, (200, 4000), (200, 4000), valid_extensions)
        # if not success:
        #     # delete_temp_files(save_path)
        #     return None, "图片压缩失败" + new_path

        return save_path, None
    except Exception as e:
        return None, f"处理图片异常: {str(e)}"


# --------------------------
# 通用工具函数
# --------------------------
def delete_temp_files(*file_paths):
    print("删除", file_paths)
    """删除临时文件（忽略不存在的文件）"""
    for path in file_paths:
        if path and os.path.exists(path):
            try:
                os.remove(path)
                print(f"已清理临时文件: {path}")
            except Exception as e:
                print(f"清理临时文件失败 {path}: {str(e)}")


# --------------------------
# 视频生成核心逻辑
# --------------------------
def ViduGenerator_pixverse(body):
    """统一处理视频生成任务（入口函数）"""
    # start_time = sr_utils.getTimes(local_time)

    task_id = body.get('taskId', None)
    queue_result = body.get('QueueResult')
    storage_service = body.get('storageService', 's3')
    bucket_name = body.get('bucketName')
    s3_storage = body.get('storedPrefix')
    model_type = body.get('modelType')
    module_name = body.get('moduleName')  #模型版本
    negative_prompt = body.get('negative_prompt')  #反向提示词
    motion_mode = body.get('mode', "normal")  # 模式
    bgm = body.get('bgm', False)  # 模式
    if bgm == 1:
        bgm = True
    else:
        bgm = False
    template_id = int(body.get('template', 0))  # 特效模板id
    print(template_id)
    camera_movement = body.get('camera_movement', None)
    style = body.get('model_style', None)
    # 必传参数校验
    required_params = ['prompt']
    for param in required_params:
        if not body.get(param) and template_id == 0:
            err_msg = f"缺少参数: {param}"
            # send_result_to_queue(queue_result, body, 0, err_msg, start_time=start_time, local_time=local_time)
            return {"status": 2, "msg": "缺少参数", "code": 6003}  # False, err_msg

    prompt = body.get('prompt', None)
    duration = int(body.get('duration', None))
    resolution = body.get('mResolution', None)
    ratio = body.get('ratio', None)
    watermark = body.get('watermark', False)
    seed = body.get('seed', 0)

    # 临时文件路径（统一管理）
    first_image_path = None
    last_image_path = None
    local_video = None
    video_result = {
        'save_path': None,
        'width': 0,
        'height': 0,
        'duration': 0,
        'fps': 0
    }
    # model_type 1文生视频 2图生视频 3首尾帧 4模板生成
    # 在整个任务过程中使用相同的trace_id
    trace_id = str(uuid.uuid4())
    try:
        # --------------------------
        # 1. 根据模型类型准备参数
        # --------------------------
        if model_type == 1:
            # 文生视频：无需图片
            create_params = {
                "aspect_ratio": ratio,
                "duration": duration,
                "model": module_name,
                "prompt": prompt,
                "quality": resolution,
                "seed": seed,
                "negative_prompt": negative_prompt,
                "motion_mode": motion_mode,
            }
            if style is None:
                if style == "auto":
                    target_list = ["anime", "3d_animation", "day", "cyberpunk", "comic"]
                    style = random.choice(target_list)
                create_params["style"] = style
            if camera_movement is None and (module_name == "v4" or module_name == "4.5"):
                create_params["camera_movement"] = camera_movement
            if ratio == None or duration == None or prompt == None or resolution == None:
                return {"status": 2, "msg": "参数不正确", "code": 6003}

            if template_id > 0:
                create_params["model"] = "v5"
                create_params["template_id"] = template_id
                create_params["sound_effect_switch"] = bgm
            # # --------------------------
            # 2. 生成视频
            # --------------------------
            print(f"开始生成视频，任务ID: {task_id}，模型类型: {model_type}")
            local_video = create_task(create_params, TEXT_TO_VIDEO_URL, trace_id)

        elif model_type == 2:
            # 文+图生视频：需要首帧图片
            first_image_url = body.get('first_image')
            if not first_image_url:
                return {"status": 2, "msg": "model_type=2需提供first_image参数", "code": 6003}  # False, "model_type=2需提供first_image参数"

            # 下载首帧图片
            first_image_path, err = download_single_image(
                first_image_url, storage_service, bucket_name
            )
            if not first_image_path:
                return {"status": 2, "msg": "首帧图片处理失败:{err}", "code": 6005}  # False, f"首帧图片处理失败: {err}"

            # 远程接口URL
            # 文件上传
            res = upload_file(first_image_path, trace_id)
            if res["ErrCode"] == 0:
                img_id = res["Resp"]["img_id"]
                create_params = {
                    "duration": duration,
                    "img_id": img_id,
                    "model": module_name,
                    "quality": resolution,
                    "seed": seed,
                    "motion_mode": motion_mode,
                }
                
                # 只有当prompt存在且非空时才添加到参数中
                if prompt is not None and prompt != "":
                    create_params["prompt"] = prompt
                    
                if negative_prompt is not None and negative_prompt != "":
                    create_params["negative_prompt"] = negative_prompt
                    
                if style is not None and style != "":
                    if style == "auto":
                        target_list = ["anime", "3d_animation", "day", "cyberpunk", "comic"]
                        style = random.choice(target_list)
                    create_params["style"] = style
                    
                if camera_movement is not None and camera_movement != "" and (module_name == "v4" or module_name == "4.5"):
                    create_params["camera_movement"] = camera_movement
                    
                if img_id == None or duration == None or resolution == None:
                    return {"status": 2, "msg": "参数不正确", "code": 6003}
                    
                if template_id > 0:
                    create_params["template_id"] = template_id
                    create_params["model"] = "v5"
                    create_params["img_id"] = img_id
                    create_params["sound_effect_switch"] = bgm
                    
                local_video = create_task(create_params, IMG_TO_VIDEO_URL, trace_id)
            else:
                return {"status": 2, "msg": ""}

        elif model_type == 3:
            api_url = FRAME_TO_VIDEO_URL
            # 首尾帧生视频：需要首帧和尾帧图片
            first_image_url = body.get('first_image')
            last_image_url = body.get('last_image')
            if not (first_image_url and last_image_url):
                return {"status": 2, "msg": "model_type=3需提供first_image和last_image参数", "code": 6003}  # False, "model_type=3需提供first_image和last_image参数"

            # 下载首帧和尾帧图片
            first_image_path, err = download_single_image(
                first_image_url, storage_service, bucket_name
            )
            if not first_image_path:
                return {"status": 2, "msg": "首帧图片处理失败:{err}", "code": 6005}  # False, f"首帧图片处理失败: {err}"
            # 远程接口URL
            # 文件上传
            first_frame_res = upload_file(first_image_path, trace_id)
            if first_frame_res["ErrCode"] != 0:
                return {"status": 2, "msg": first_frame_res["ErrMsg"], "code": 6001}

            last_image_path, err = download_single_image(
                last_image_url, storage_service, bucket_name
            )
            if not last_image_path:
                return {"status": 2, "msg": "尾帧图片处理失败:{err}", "code": 6005}  # False, f"尾帧图片处理失败: {err}"

            last_frame_res = upload_file(last_image_path, trace_id)
            if last_frame_res["ErrCode"] != 0:
                return {"status": 2, "msg": last_frame_res["ErrMsg"], "code": 6005}

            create_params = {
                "model": module_name,
                "duration": duration,
                "quality": resolution,
                "seed": seed,
                "first_frame_img": first_frame_res["Resp"]["img_id"],
                "last_frame_img": last_frame_res["Resp"]["img_id"],
                "motion_mode": motion_mode,
            }

            # 只有当prompt存在且非空时才添加到参数中
            if prompt is not None and prompt != "":
                create_params["prompt"] = prompt
                
            if negative_prompt is not None and negative_prompt != "":
                create_params["negative_prompt"] = negative_prompt
                
            if style is not None and style != "":
                if style == "auto":
                    target_list = ["anime", "3d_animation", "day", "cyberpunk", "comic"]
                    style = random.choice(target_list)
                create_params["style"] = style
                
            if camera_movement is not None and camera_movement != "" and (module_name == "v4" or module_name == "4.5"):
                create_params["camera_movement"] = camera_movement
                
            if last_frame_res["Resp"]["img_id"] == None or first_frame_res["Resp"]["img_id"] == None or duration == None or resolution == None:
                return {"status": 2, "msg": "参数不正确", "code": 6003}
                
            if template_id > 0:
                api_url = IMG_TO_VIDEO_URL
                # 双图逻辑
                create_params["model"] = "v4.5"
                create_params["template_id"] = template_id
                create_params["sound_effect_switch"] = bgm
                # 根据PixVerse API文档，使用模板时需要移除首尾帧参数，并使用img_ids参数
                create_params.pop('first_frame_img', None)
                create_params.pop('last_frame_img', None)
                create_params['img_ids'] = [first_frame_res["Resp"]["img_id"], last_frame_res["Resp"]["img_id"]]
                
            # create_params = {'prompt': '让两个人跑起来，最后亲吻', 'model': 'v3.5', 'duration': 5, 'quality': '360p', 'seed': 0, 'first_frame_img': 667035698, 'last_frame_img': 667035854}
            local_video = create_task(create_params, api_url, trace_id)
        else:
            return False, f"不支持的模型类型: {model_type}"

        if local_video["ErrCode"] != 0:
            return {"status": 2, "msg": local_video["ErrMsg"], "code": 6003}
        else:
            if local_video["Resp"]["status"] == 7:  # 审核未通过
                return {"status": 7, "msg": "nsfw", "code": 7202}
            elif local_video["Resp"]["status"] == 1:  # 生成成功
                # 3. 下载并保存视频到本地
                # 生成带任务ID的文件名（便于追踪）
                filename = f"video_task_{task_id}.mp4" if task_id else None
                video_result = download_video(
                    video_url=local_video["Resp"]["url"],
                    save_dir="./download",
                    filename=filename
                )
                # --------------------------
                # 3. 上传视频到S3
                # --------------------------
                print(f"开始上传视频到S3: {video_result['save_path']}")
                s3_video_key, err = upload_video_to_s3(
                    video_result['save_path'], s3_storage, 'r2', bucket_name
                )
                if not s3_video_key:
                    return False, f"视频上传r2失败: {err}"
                data = {
                    "status": 1,
                    "width": local_video["Resp"]["outputWidth"],
                    "height": local_video["Resp"]["outputHeight"],
                    'url': s3_video_key,
                    "model_mode": motion_mode,  # 模型类型标准模式,
                    "local_path": video_result['save_path'],
                    'msg': "success",
                    "code": ""
                }
                return data
            else:  # 生成失败
                return {"status": 2, "msg": "", "code": 6001}

    except Exception as e:
        err_msg = f"任务处理异常: {str(e)}"
        print(f"任务{task_id}异常: {traceback.format_exc()}")

        return {"status": 2, "msg": err_msg, "code": 6001}  # False, err_msg



# 使用示例
if __name__ == "__main__":
    task_params ={'taskId': '1761099135693XVPuTkdEFCzk0nyS', 'prompt': '', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1761098473838.png', 'last_image': '', 'duration': '5', 'resolution': '720p', 'ratio': 'adaptive', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 2, 'moduleType': 22, 'mResolution': '540p', 'template': '360749956342336', 'p_t_b': 0, 'beast': '', 'bgm': 1}
    res = ViduGenerator_pixverse(task_params)
    print("制作结果", res)
